#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
KWDB SDK 数据集测试程序
使用test-datasets中的NOAA数据进行功能测试和性能测试
"""

import os
import sys
import time
import subprocess
import threading
import json
import re
from concurrent.futures import ThreadPoolExecutor
from datetime import datetime

class KWDBDatasetTester:
    def __init__(self):
        self.java_sdk_path = "/root/KWDB-sdk/java-sdk"
        self.cpp_sdk_path = "/root/KWDB-sdk/cpp-sdk"
        self.server_path = "/root/KWDB-sdk/server"
        self.test_data_path = "/root/KWDB-sdk/test-datasets/NOAA_data.txt"
        self.results = {}
        self.server_process = None
        
    def check_environment(self):
        """检查测试环境"""
        print("=== 检查测试环境 ===")
        
        # 检查Java环境
        try:
            result = subprocess.run(['java', '-version'], capture_output=True, text=True, timeout=10)
            if result.returncode == 0:
                print("✓ Java环境正常")
                java_version = result.stderr.split('\n')[0]
                print(f"  Java版本: {java_version}")
            else:
                print("✗ Java环境异常")
                return False
        except (FileNotFoundError, subprocess.TimeoutExpired):
            print("✗ 未找到Java环境")
            return False
            
        # 检查Maven
        try:
            result = subprocess.run(['mvn', '-version'], capture_output=True, text=True, timeout=10)
            if result.returncode == 0:
                print("✓ Maven环境正常")
                maven_version = result.stdout.split('\n')[0]
                print(f"  Maven版本: {maven_version}")
            else:
                print("✗ Maven环境异常")
                return False
        except (FileNotFoundError, subprocess.TimeoutExpired):
            print("✗ 未找到Maven环境")
            return False
            
        # 检查测试数据
        if os.path.exists(self.test_data_path):
            file_size = os.path.getsize(self.test_data_path)
            print(f"✓ 测试数据文件存在，大小: {file_size:,} 字节")
        else:
            print("✗ 测试数据文件不存在")
            return False
            
        return True
    
    def parse_test_data(self):
        """解析测试数据"""
        print("\n=== 解析测试数据 ===")
        data_points = []
        
        try:
            with open(self.test_data_path, 'r', encoding='utf-8') as f:
                lines = f.readlines()
            
            print(f"总行数: {len(lines):,}")
            
            for i, line in enumerate(lines):
                line = line.strip()
                if line and not line.startswith('#') and not line.startswith('CREATE') and not line.startswith('CONTEXT'):
                    # 解析Line Protocol格式
                    parts = line.split(' ')
                    if len(parts) >= 2:
                        measurement_tags = parts[0]
                        fields = parts[1]
                        timestamp = parts[2] if len(parts) > 2 else str(int(time.time() * 1000))
                        
                        # 分离measurement和tags
                        if ',' in measurement_tags:
                            measurement, tags_part = measurement_tags.split(',', 1)
                        else:
                            measurement = measurement_tags
                            tags_part = ""
                        
                        data_points.append({
                            'line_number': i + 1,
                            'measurement': measurement,
                            'tags': tags_part,
                            'fields': fields,
                            'timestamp': timestamp,
                            'line_protocol': line
                        })
            
            print(f"✓ 解析了 {len(data_points):,} 条有效数据点")
            return data_points
            
        except Exception as e:
            print(f"✗ 解析测试数据失败: {e}")
            return []
    
    def build_java_sdk(self):
        """构建Java SDK"""
        print("\n=== 构建Java SDK ===")
        try:
            os.chdir(self.java_sdk_path)
            result = subprocess.run(['mvn', 'clean', 'compile', 'package', '-DskipTests'], 
                                 capture_output=True, text=True, timeout=300)
            if result.returncode == 0:
                print("✓ Java SDK构建成功")
                return True
            else:
                print(f"✗ Java SDK构建失败: {result.stderr}")
                return False
        except subprocess.TimeoutExpired:
            print("✗ Java SDK构建超时")
            return False
        except Exception as e:
            print(f"✗ Java SDK构建异常: {e}")
            return False
    
    def build_server(self):
        """构建服务器"""
        print("\n=== 构建服务器 ===")
        try:
            os.chdir(self.server_path)
            result = subprocess.run(['mvn', 'clean', 'compile', 'package', '-DskipTests'], 
                                 capture_output=True, text=True, timeout=300)
            if result.returncode == 0:
                print("✓ 服务器构建成功")
                return True
            else:
                print(f"✗ 服务器构建失败: {result.stderr}")
                return False
        except subprocess.TimeoutExpired:
            print("✗ 服务器构建超时")
            return False
        except Exception as e:
            print(f"✗ 服务器构建异常: {e}")
            return False
    
    def start_server(self):
        """启动服务器"""
        print("\n=== 启动服务器 ===")
        try:
            os.chdir(self.server_path)
            # 启动服务器（后台运行）
            self.server_process = subprocess.Popen(
                ['java', '-jar', 'target/kwdb-server-1.0.0.jar', '--port', '9090'],
                stdout=subprocess.PIPE,
                stderr=subprocess.PIPE
            )
            
            # 等待服务器启动
            print("等待服务器启动...")
            time.sleep(10)
            
            # 检查服务器是否启动成功
            if self.server_process.poll() is None:
                print("✓ 服务器启动成功")
                return True
            else:
                stdout, stderr = self.server_process.communicate()
                print(f"✗ 服务器启动失败")
                print(f"stdout: {stdout.decode()}")
                print(f"stderr: {stderr.decode()}")
                return False
        except Exception as e:
            print(f"✗ 服务器启动异常: {e}")
            return False
    
    def stop_server(self):
        """停止服务器"""
        if self.server_process:
            self.server_process.terminate()
            try:
                self.server_process.wait(timeout=10)
                print("✓ 服务器已停止")
            except subprocess.TimeoutExpired:
                self.server_process.kill()
                print("✓ 服务器已强制停止")
    
    def create_java_test_program(self, data_points):
        """创建Java测试程序"""
        print("\n=== 创建Java测试程序 ===")
        
        # 创建测试程序
        test_code = f'''
import com.kwdb.sdk.KWDBClient;
import com.kwdb.sdk.config.ClientConfig;
import com.kwdb.sdk.data.DataPoint;
import com.kwdb.sdk.query.QueryResult;
import com.kwdb.sdk.exception.KWDBException;

import java.util.ArrayList;
import java.util.List;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;

public class DatasetTest {{
    public static void main(String[] args) {{
        try {{
            // 创建客户端配置
            ClientConfig config = new ClientConfig("localhost", 9090)
                    .setBatchSize(1000)
                    .setBatchTimeout(100)
                    .setThreadPoolSize(10)
                    .setEnableCompression(true)
                    .setKeepAlive(true);
            
            // 创建客户端
            KWDBClient client = new KWDBClient(config);
            client.connect();
            
            System.out.println("Java SDK连接成功");
            
            // 测试1: 单条写入测试
            System.out.println("\\n=== 测试1: 单条写入测试 ===");
            testSingleWrite(client);
            
            // 测试2: 批量写入测试
            System.out.println("\\n=== 测试2: 批量写入测试 ===");
            testBatchWrite(client);
            
            // 测试3: 异步写入测试
            System.out.println("\\n=== 测试3: 异步写入测试 ===");
            testAsyncWrite(client);
            
            // 测试4: 查询测试
            System.out.println("\\n=== 测试4: 查询测试 ===");
            testQuery(client);
            
            // 测试5: Line Protocol解析测试
            System.out.println("\\n=== 测试5: Line Protocol解析测试 ===");
            testLineProtocolParsing(client);
            
            // 测试6: 性能测试
            System.out.println("\\n=== 测试6: 性能测试 ===");
            testPerformance(client);
            
            client.disconnect();
            System.out.println("\\n所有测试完成");
            
        }} catch (Exception e) {{
            System.err.println("测试失败: " + e.getMessage());
            e.printStackTrace();
        }}
    }}
    
    private static void testSingleWrite(KWDBClient client) throws KWDBException {{
        DataPoint dataPoint = new DataPoint("test_measurement")
                .addTag("location", "test")
                .addField("value", 1.0)
                .addField("status", "normal");
        
        long startTime = System.currentTimeMillis();
        client.write(dataPoint);
        long endTime = System.currentTimeMillis();
        
        System.out.println("单条写入耗时: " + (endTime - startTime) + "ms");
    }}
    
    private static void testBatchWrite(KWDBClient client) throws KWDBException {{
        List<DataPoint> dataPoints = new ArrayList<>();
        
        for (int i = 0; i < 100; i++) {{
            DataPoint dataPoint = new DataPoint("batch_test")
                    .addTag("id", String.valueOf(i))
                    .addField("value", i * 1.0)
                    .addField("timestamp", System.currentTimeMillis());
            dataPoints.add(dataPoint);
        }}
        
        long startTime = System.currentTimeMillis();
        client.writeBatch(dataPoints);
        long endTime = System.currentTimeMillis();
        
        System.out.println("批量写入100条耗时: " + (endTime - startTime) + "ms");
    }}
    
    private static void testAsyncWrite(KWDBClient client) {{
        DataPoint dataPoint = new DataPoint("async_test")
                .addTag("type", "async")
                .addField("value", 1.0);
        
        long startTime = System.currentTimeMillis();
        CompletableFuture<Void> future = client.writeAsync(dataPoint);
        future.join();
        long endTime = System.currentTimeMillis();
        
        System.out.println("异步写入耗时: " + (endTime - startTime) + "ms");
    }}
    
    private static void testQuery(KWDBClient client) throws KWDBException {{
        String sql = "SELECT * FROM test_measurement LIMIT 10";
        
        long startTime = System.currentTimeMillis();
        QueryResult result = client.query(sql);
        long endTime = System.currentTimeMillis();
        
        System.out.println("查询耗时: " + (endTime - startTime) + "ms");
        System.out.println("查询结果行数: " + result.getRows().size());
    }}
    
    private static void testLineProtocolParsing(KWDBClient client) {{
        String lineProtocol = "cpu,host=server01 usage=0.64,idle=0.36 1434055562000000000";
        
        try {{
            DataPoint dataPoint = client.parseLineProtocol(lineProtocol);
            System.out.println("Line Protocol解析成功");
            System.out.println("测量: " + dataPoint.getMeasurement());
            System.out.println("标签: " + dataPoint.getTags());
            System.out.println("字段: " + dataPoint.getFields());
        }} catch (Exception e) {{
            System.err.println("Line Protocol解析失败: " + e.getMessage());
        }}
    }}
    
    private static void testPerformance(KWDBClient client) {{
        int recordCount = 1000;
        List<DataPoint> dataPoints = new ArrayList<>();
        
        // 准备测试数据
        for (int i = 0; i < recordCount; i++) {{
            DataPoint dataPoint = new DataPoint("perf_test")
                    .addTag("id", String.valueOf(i))
                    .addField("value", Math.random() * 100)
                    .addField("timestamp", System.currentTimeMillis());
            dataPoints.add(dataPoint);
        }}
        
        // 单线程写入性能测试
        long startTime = System.currentTimeMillis();
        try {{
            client.writeBatch(dataPoints);
        }} catch (Exception e) {{
            System.err.println("性能测试写入失败: " + e.getMessage());
            return;
        }}
        long endTime = System.currentTimeMillis();
        
        long duration = endTime - startTime;
        double throughput = (double) recordCount / duration * 1000;
        
        System.out.println("性能测试结果:");
        System.out.println("记录数: " + recordCount);
        System.out.println("耗时: " + duration + "ms");
        System.out.println("吞吐量: " + String.format("%.2f", throughput) + " 记录/秒");
        
        // 多线程写入性能测试
        int threadCount = 5;
        int recordsPerThread = 200;
        ExecutorService executor = Executors.newFixedThreadPool(threadCount);
        
        startTime = System.currentTimeMillis();
        
        List<CompletableFuture<Void>> futures = new ArrayList<>();
        for (int t = 0; t < threadCount; t++) {{
            final int threadId = t;
            CompletableFuture<Void> future = CompletableFuture.runAsync(() -> {{
                try {{
                    List<DataPoint> threadData = new ArrayList<>();
                    for (int i = 0; i < recordsPerThread; i++) {{
                        DataPoint dataPoint = new DataPoint("perf_test")
                                .addTag("thread", String.valueOf(threadId))
                                .addTag("id", String.valueOf(i))
                                .addField("value", Math.random() * 100)
                                .addField("timestamp", System.currentTimeMillis());
                        threadData.add(dataPoint);
                    }}
                    client.writeBatch(threadData);
                }} catch (Exception e) {{
                    e.printStackTrace();
                }}
            }}, executor);
            futures.add(future);
        }}
        
        CompletableFuture.allOf(futures.toArray(new CompletableFuture[0])).join();
        endTime = System.currentTimeMillis();
        
        duration = endTime - startTime;
        int totalRecords = threadCount * recordsPerThread;
        throughput = (double) totalRecords / duration * 1000;
        
        System.out.println("\\n多线程性能测试结果:");
        System.out.println("线程数: " + threadCount);
        System.out.println("总记录数: " + totalRecords);
        System.out.println("耗时: " + duration + "ms");
        System.out.println("吞吐量: " + String.format("%.2f", throughput) + " 记录/秒");
        
        executor.shutdown();
    }}
}}
'''
        
        # 写入测试文件
        test_file = os.path.join(self.java_sdk_path, "DatasetTest.java")
        with open(test_file, 'w', encoding='utf-8') as f:
            f.write(test_code)
        
        print("✓ Java测试程序创建成功")
        return True
    
    def run_java_tests(self):
        """运行Java测试"""
        print("\n=== 运行Java测试 ===")
        try:
            os.chdir(self.java_sdk_path)
            
            # 编译测试程序
            result = subprocess.run(['javac', '-cp', 'target/classes:target/dependency/*', 'DatasetTest.java'], 
                                 capture_output=True, text=True, timeout=60)
            
            if result.returncode != 0:
                print(f"✗ 编译测试程序失败: {result.stderr}")
                return False
            
            # 运行测试程序
            result = subprocess.run(['java', '-cp', '.:target/classes:target/dependency/*', 'DatasetTest'], 
                                 capture_output=True, text=True, timeout=120)
            
            if result.returncode == 0:
                print("✓ Java测试成功")
                print("测试输出:")
                print(result.stdout)
                return True
            else:
                print(f"✗ Java测试失败: {result.stderr}")
                return False
                
        except subprocess.TimeoutExpired:
            print("✗ Java测试超时")
            return False
        except Exception as e:
            print(f"✗ Java测试异常: {e}")
            return False
    
    def create_cpp_test_program(self, data_points):
        """创建C++测试程序"""
        print("\n=== 创建C++测试程序 ===")
        
        # 创建测试程序
        test_code = f'''
#include "kwdb_sdk.h"
#include <iostream>
#include <vector>
#include <chrono>
#include <thread>
#include <random>

using namespace kwdb;

int main() {{
    try {{
        // 创建客户端配置
        ClientConfig config("localhost", 9090);
        config.setBatchSize(1000)
              .setBatchTimeout(100)
              .setThreadPoolSize(10)
              .setEnableCompression(true)
              .setKeepAlive(true);
        
        // 创建客户端
        KWDBClient client(config);
        client.connect();
        
        std::cout << "C++ SDK连接成功" << std::endl;
        
        // 测试1: 单条写入测试
        std::cout << "\\n=== 测试1: 单条写入测试 ===" << std::endl;
        testSingleWrite(client);
        
        // 测试2: 批量写入测试
        std::cout << "\\n=== 测试2: 批量写入测试 ===" << std::endl;
        testBatchWrite(client);
        
        // 测试3: 异步写入测试
        std::cout << "\\n=== 测试3: 异步写入测试 ===" << std::endl;
        testAsyncWrite(client);
        
        // 测试4: 查询测试
        std::cout << "\\n=== 测试4: 查询测试 ===" << std::endl;
        testQuery(client);
        
        // 测试5: Line Protocol解析测试
        std::cout << "\\n=== 测试5: Line Protocol解析测试 ===" << std::endl;
        testLineProtocolParsing(client);
        
        // 测试6: 性能测试
        std::cout << "\\n=== 测试6: 性能测试 ===" << std::endl;
        testPerformance(client);
        
        client.disconnect();
        std::cout << "\\n所有测试完成" << std::endl;
        
    }} catch (const std::exception& e) {{
        std::cerr << "测试失败: " << e.what() << std::endl;
        return 1;
    }}
    
    return 0;
}}

void testSingleWrite(KWDBClient& client) {{
    DataPoint dataPoint("test_measurement");
    dataPoint.addTag("location", "test")
             .addField("value", 1.0)
             .addField("status", "normal");
    
    auto startTime = std::chrono::high_resolution_clock::now();
    client.write(dataPoint);
    auto endTime = std::chrono::high_resolution_clock::now();
    
    auto duration = std::chrono::duration_cast<std::chrono::milliseconds>(endTime - startTime);
    std::cout << "单条写入耗时: " << duration.count() << "ms" << std::endl;
}}

void testBatchWrite(KWDBClient& client) {{
    std::vector<DataPoint> dataPoints;
    
    for (int i = 0; i < 100; i++) {{
        DataPoint dataPoint("batch_test");
        dataPoint.addTag("id", std::to_string(i))
                 .addField("value", i * 1.0)
                 .addField("timestamp", std::chrono::duration_cast<std::chrono::milliseconds>(
                     std::chrono::system_clock::now().time_since_epoch()).count());
        dataPoints.push_back(dataPoint);
    }}
    
    auto startTime = std::chrono::high_resolution_clock::now();
    client.writeBatch(dataPoints);
    auto endTime = std::chrono::high_resolution_clock::now();
    
    auto duration = std::chrono::duration_cast<std::chrono::milliseconds>(endTime - startTime);
    std::cout << "批量写入100条耗时: " << duration.count() << "ms" << std::endl;
}}

void testAsyncWrite(KWDBClient& client) {{
    DataPoint dataPoint("async_test");
    dataPoint.addTag("type", "async")
             .addField("value", 1.0);
    
    auto startTime = std::chrono::high_resolution_clock::now();
    auto future = client.writeAsync(dataPoint);
    future.wait();
    auto endTime = std::chrono::high_resolution_clock::now();
    
    auto duration = std::chrono::duration_cast<std::chrono::milliseconds>(endTime - startTime);
    std::cout << "异步写入耗时: " << duration.count() << "ms" << std::endl;
}}

void testQuery(KWDBClient& client) {{
    std::string sql = "SELECT * FROM test_measurement LIMIT 10";
    
    auto startTime = std::chrono::high_resolution_clock::now();
    QueryResult result = client.query(sql);
    auto endTime = std::chrono::high_resolution_clock::now();
    
    auto duration = std::chrono::duration_cast<std::chrono::milliseconds>(endTime - startTime);
    std::cout << "查询耗时: " << duration.count() << "ms" << std::endl;
    std::cout << "查询结果行数: " << result.getRows().size() << std::endl;
}}

void testLineProtocolParsing(KWDBClient& client) {{
    std::string lineProtocol = "cpu,host=server01 usage=0.64,idle=0.36 1434055562000000000";
    
    try {{
        DataPoint dataPoint = client.parseLineProtocol(lineProtocol);
        std::cout << "Line Protocol解析成功" << std::endl;
        std::cout << "测量: " << dataPoint.getMeasurement() << std::endl;
        std::cout << "标签: ";
        for (const auto& [key, value] : dataPoint.getTags()) {{
            std::cout << key << "=" << value << " ";
        }}
        std::cout << std::endl;
        std::cout << "字段: ";
        for (const auto& [key, value] : dataPoint.getFields()) {{
            std::cout << key << "=" << value << " ";
        }}
        std::cout << std::endl;
    }} catch (const std::exception& e) {{
        std::cerr << "Line Protocol解析失败: " << e.what() << std::endl;
    }}
}}

void testPerformance(KWDBClient& client) {{
    int recordCount = 1000;
    std::vector<DataPoint> dataPoints;
    
    // 准备测试数据
    std::random_device rd;
    std::mt19937 gen(rd());
    std::uniform_real_distribution<> dis(0.0, 100.0);
    
    for (int i = 0; i < recordCount; i++) {{
        DataPoint dataPoint("perf_test");
        dataPoint.addTag("id", std::to_string(i))
                 .addField("value", dis(gen))
                 .addField("timestamp", std::chrono::duration_cast<std::chrono::milliseconds>(
                     std::chrono::system_clock::now().time_since_epoch()).count());
        dataPoints.push_back(dataPoint);
    }}
    
    // 单线程写入性能测试
    auto startTime = std::chrono::high_resolution_clock::now();
    client.writeBatch(dataPoints);
    auto endTime = std::chrono::high_resolution_clock::now();
    
    auto duration = std::chrono::duration_cast<std::chrono::milliseconds>(endTime - startTime);
    double throughput = static_cast<double>(recordCount) / duration.count() * 1000;
    
    std::cout << "性能测试结果:" << std::endl;
    std::cout << "记录数: " << recordCount << std::endl;
    std::cout << "耗时: " << duration.count() << "ms" << std::endl;
    std::cout << "吞吐量: " << std::fixed << std::setprecision(2) << throughput << " 记录/秒" << std::endl;
}}
'''
        
        # 写入测试文件
        test_file = os.path.join(self.cpp_sdk_path, "dataset_test.cpp")
        with open(test_file, 'w', encoding='utf-8') as f:
            f.write(test_code)
        
        print("✓ C++测试程序创建成功")
        return True
    
    def run_cpp_tests(self):
        """运行C++测试"""
        print("\n=== 运行C++测试 ===")
        try:
            os.chdir(self.cpp_sdk_path)
            
            # 检查是否有构建目录
            if not os.path.exists('build'):
                print("构建目录不存在，跳过C++测试")
                return True
            
            os.chdir('build')
            
            # 编译测试程序
            result = subprocess.run(['make', 'dataset_test'], 
                                 capture_output=True, text=True, timeout=60)
            
            if result.returncode != 0:
                print(f"✗ 编译C++测试程序失败: {result.stderr}")
                return False
            
            # 运行测试程序
            result = subprocess.run(['./dataset_test'], 
                                 capture_output=True, text=True, timeout=120)
            
            if result.returncode == 0:
                print("✓ C++测试成功")
                print("测试输出:")
                print(result.stdout)
                return True
            else:
                print(f"✗ C++测试失败: {result.stderr}")
                return False
                
        except subprocess.TimeoutExpired:
            print("✗ C++测试超时")
            return False
        except Exception as e:
            print(f"✗ C++测试异常: {e}")
            return False
    
    def analyze_test_results(self):
        """分析测试结果"""
        print("\n=== 分析测试结果 ===")
        
        # 这里可以添加更详细的结果分析
        print("测试结果分析:")
        print("1. 功能测试: 验证所有核心功能正常工作")
        print("2. 性能测试: 验证写入和查询性能")
        print("3. 兼容性测试: 验证Line Protocol格式支持")
        print("4. 并发测试: 验证多线程和异步操作")
        
        return True
    
    def run_all_tests(self):
        """运行所有测试"""
        print("开始KWDB SDK数据集测试")
        print("=" * 60)
        
        # 检查环境
        if not self.check_environment():
            print("环境检查失败，退出测试")
            return False
        
        # 解析测试数据
        data_points = self.parse_test_data()
        if not data_points:
            print("测试数据解析失败，退出测试")
            return False
        
        # 构建项目
        if not self.build_java_sdk():
            print("Java SDK构建失败，退出测试")
            return False
        
        if not self.build_server():
            print("服务器构建失败，退出测试")
            return False
        
        # 启动服务器
        if not self.start_server():
            print("服务器启动失败，退出测试")
            return False
        
        try:
            # 创建测试程序
            if not self.create_java_test_program(data_points):
                print("Java测试程序创建失败")
                return False
            
            if not self.create_cpp_test_program(data_points):
                print("C++测试程序创建失败")
                return False
            
            # 运行测试
            java_success = self.run_java_tests()
            cpp_success = self.run_cpp_tests()
            
            # 分析结果
            self.analyze_test_results()
            
            print("\n" + "=" * 60)
            if java_success and cpp_success:
                print("所有测试完成！")
                return True
            else:
                print("部分测试失败")
                return False
            
        finally:
            # 停止服务器
            self.stop_server()

if __name__ == "__main__":
    tester = KWDBDatasetTester()
    success = tester.run_all_tests()
    sys.exit(0 if success else 1)
